import psutil
from nonebot import on_regex
from nonebot.adapters.onebot.v11 import Message, GroupMessageEvent, Bot
from nonebot.exception import FinishedException  # 引入 FinishedException

# 监听消息中是否包含"disk"，不区分大小写
disk_command = on_regex(pattern=r'(?i)^disk$', priority=1)

@disk_command.handle()
async def handle_disk(bot: Bot, event: GroupMessageEvent):
    try:
        # 获取磁盘分区信息
        partitions = psutil.disk_partitions()

        # 存储每个磁盘的使用情况
        disk_info = ""

        for partition in partitions:
            if "loop" in partition.device:
                continue
            
            # 获取每个分区的使用情况
            usage = psutil.disk_usage(partition.mountpoint)
            total_disk_gb = usage.total / (1024 ** 3)  # 转换为GB
            used_disk_gb = usage.used / (1024 ** 3)    # 转换为GB
            disk_usage_percent = usage.percent

            # 获取磁盘类型
            # disk_type = "固态硬盘" if "ssd" in partition.device.lower() else "机械硬盘"

            # 构建每个磁盘的信息
            disk_info += (
                # f"{partition.device} ({disk_type}) - 已用: {used_disk_gb:.2f} GB / 总空间: {total_disk_gb:.2f} GB，"
                f"{partition.device} - 已用: {used_disk_gb:.2f} GB / 总空间: {total_disk_gb:.2f} GB，"
                f"使用率: {disk_usage_percent:.2f}%\n"
            )

        if not disk_info:
            disk_info = "未能获取到磁盘信息"

        # 输出磁盘占用信息
        await disk_command.finish(message=Message(disk_info))

    except Exception as e:
        # 判断异常类型
        if not isinstance(e, FinishedException):
            print(f"发生错误: {e}")
            await disk_command.finish(message=Message(f"发生错误: {e}"))
